#define STB_IMAGE_IMPLEMENTATION
#include "./stb/stb_image.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "./stb/stb_image_write.h"
#define STB_IMAGE_RESIZE_IMPLEMENTATION
#include "./stb/stb_image_resize.h"
#include "./stb/drawing.hpp"
#include "./stb/common_data_struct.h"
#include <stdint.h>  // 确保 uint8_t 定义
#include <stdio.h>
#include <string>


// 图像基本信息结构体 (兼容原有代码)
typedef struct {
    uint8_t* data;   // 像素数据指针 (RGBA格式)
    int width;       // 图像宽度
    int height;      // 图像高度
    int channels; 
} FrameInfo;

typedef struct {
    float x0, y0, x1, y1;
} Box_t;

static const uint8_t color[3] = {0xff, 0, 0};

int fill_dim(uint8_t* data, int h, int w);
void read_img(FrameInfo& in, Img& img, const std::string& img_path);


uint8_t* rgb_to_rgba(const uint8_t* src, int width, int height, uint8_t alpha) {
    if (!src || width <= 0 || height <= 0) return NULL;
    
    uint8_t* dst = (uint8_t*)malloc(width * height * 4);
    if (!dst) return NULL;

    for (int i = 0; i < width * height; i++) {
        dst[i*4]   = src[i*3];     // R
        dst[i*4+1] = src[i*3+1];   // G
        dst[i*4+2] = src[i*3+2];   // B
        dst[i*4+3] = alpha;        // A
    }
    return dst;
}

uint8_t* rgba_to_rgb(const uint8_t* src, int width, int height) {
    if (!src || width <= 0 || height <= 0) return NULL;
    
    uint8_t* dst = (uint8_t*)malloc(width * height * 3);
    if (!dst) return NULL;

    for (int i = 0; i < width * height; i++) {
        dst[i*3]   = src[i*4];     // R
        dst[i*3+1] = src[i*4+1];   // G
        dst[i*3+2] = src[i*4+2];   // B
        // 忽略src[i*4+3] (Alpha)
    }
    return dst;
}

// 读取图像并填充结构体
void read_img(FrameInfo& in, Img& img, const std::string& img_path) {
    printf("img_path===> %s\n", img_path.c_str());

    int ori_img_w, ori_img_h, comp;
    unsigned char* imagedata = stbi_load(img_path.c_str(), &ori_img_w, &ori_img_h, &comp, 0); // 不强制通道数
    if (!imagedata) {
        fprintf(stderr, "Error loading image: %s\n", img_path.c_str());
        return;
    }

    if (comp == 3) {
        printf("comp == 3 \n");
    }
    // 填充FrameInfo
    in.data = imagedata;
    in.width = ori_img_w;
    in.height = ori_img_h;
    in.channels = comp;
    printf("input img==> %d %d %d \n", in.width, in.height, in.channels);

    // 填充Img
    img.w = ori_img_w;
    img.h = ori_img_h;
    img.c = comp;        
    img.w_stride = ori_img_w * comp;
    img.data = imagedata;

    // 调试输出
    printf("Image loaded: width=%d, height=%d, channels=%d\n", ori_img_w, ori_img_h, comp);
}

void resize_image(const Img& src_img, Img& dst_img, int dst_width, int dst_height) {
    // 清空输出结构体
    memset(&dst_img, 0, sizeof(Img));
    
    if (!src_img.data || dst_width <= 0 || dst_height <= 0) {
        fprintf(stderr, "Invalid resize parameters\n");
        return;
    }

    // 初始化输出图像结构
    dst_img.w = dst_width;
    dst_img.h = dst_height;
    dst_img.c = src_img.c;
    dst_img.w_stride = dst_width * src_img.c;
    dst_img.data = (uint8_t*)malloc(dst_height * dst_img.w_stride);
    
    if (!dst_img.data) {
        fprintf(stderr, "Memory allocation failed for %d bytes\n", 
                dst_height * dst_img.w_stride);
        return;
    }

    // 调用stb_image_resize进行缩放
    stbir_resize(
        src_img.data, src_img.w, src_img.h, src_img.w_stride,
        dst_img.data, dst_width, dst_height, dst_img.w_stride,
        STBIR_TYPE_UINT8, src_img.c,
        (src_img.c == 4) ? STBIR_ALPHA_CHANNEL_NONE : 0,
        STBIR_FLAG_ALPHA_PREMULTIPLIED,
        STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
        STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT,
        STBIR_COLORSPACE_SRGB, NULL
    );
}

/**
 * YOLO LetterBox缩放函数
 * @param src_img 输入图像结构体
 * @param dst_img 输出图像结构体（需调用者释放data内存）
 * @param target_size 目标尺寸（正方形边长）
 * @param color 填充颜色（RGB格式，默认114的灰色）
 */
void yolo_resize_img(const Img& src_img, Img& dst_img, int target_size, const uint8_t color[3]) {
    // 清空输出结构体
    memset(&dst_img, 0, sizeof(Img));
    
    if (!src_img.data || target_size <= 0) {
        fprintf(stderr, "Invalid resize parameters\n");
        return;
    }

    // 计算缩放比例（保持宽高比）
    float scale = fminf((float)target_size / src_img.w, (float)target_size / src_img.h);
    int new_w = (int)(src_img.w * scale);
    int new_h = (int)(src_img.h * scale);

    // 计算填充位置（居中）
    int dw = target_size - new_w;
    int dh = target_size - new_h;
    int top = dh / 2;
    int bottom = dh - top;
    int left = dw / 2;
    int right = dw - left;

    // 初始化输出图像
    dst_img.w = target_size;
    dst_img.h = target_size;
    dst_img.c = src_img.c;
    dst_img.w_stride = target_size * src_img.c;
    dst_img.data = (uint8_t*)malloc(target_size * target_size * src_img.c);
    
    if (!dst_img.data) {
        fprintf(stderr, "Memory allocation failed\n");
        return;
    }

    // 创建填充画布
    for (int i = 0; i < target_size * target_size; i++) {
        for (int c = 0; c < src_img.c; c++) {
            if (c < 3) { // RGB通道
                dst_img.data[i * src_img.c + c] = color[c];
            } else if (c == 3) { // Alpha通道（如果有）
                dst_img.data[i * src_img.c + c] = 255; // 完全不透明
            }
        }
    }

    // 缩放源图像
    uint8_t* resized_data = (uint8_t*)malloc(new_w * new_h * src_img.c);
    stbir_resize(
        src_img.data, src_img.w, src_img.h, src_img.w_stride,
        resized_data, new_w, new_h, new_w * src_img.c,
        STBIR_TYPE_UINT8, src_img.c,
        (src_img.c == 4) ? STBIR_ALPHA_CHANNEL_NONE : 0,
        STBIR_FLAG_ALPHA_PREMULTIPLIED,
        STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
        STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT,
        STBIR_COLORSPACE_SRGB, NULL
    );

    // 将缩放后的图像复制到画布中央
    for (int y = 0; y < new_h; y++) {
        for (int x = 0; x < new_w; x++) {
            for (int c = 0; c < src_img.c; c++) {
                int dst_pos = ((y + top) * target_size + (x + left)) * src_img.c + c;
                int src_pos = (y * new_w + x) * src_img.c + c;
                dst_img.data[dst_pos] = resized_data[src_pos];
            }
        }
    }

    free(resized_data);
}

void crop_image(const Img& in_img, const Box_t& box, Img& out_img) {
    // 1. 边界检查（手动实现）
    int x0 = (int)box.x0;
    int y0 = (int)box.y0;
    int x1 = (int)box.x1;
    int y1 = (int)box.y1;

    // 确保坐标不越界
    x0 = (x0 < 0) ? 0 : (x0 >= in_img.w ? in_img.w - 1 : x0);
    y0 = (y0 < 0) ? 0 : (y0 >= in_img.h ? in_img.h - 1 : y0);
    x1 = (x1 < 0) ? 0 : (x1 >= in_img.w ? in_img.w - 1 : x1);
    y1 = (y1 < 0) ? 0 : (y1 >= in_img.h ? in_img.h - 1 : y1);

    // 确保 x1 > x0 且 y1 > y0
    if (x1 <= x0 || y1 <= y0) {
        out_img.w = out_img.h = out_img.c = 0;
        out_img.data = nullptr;
        printf("crop box error \n");
        return;
    }

    // 2. 计算裁剪区域
    int crop_w = x1 - x0;
    int crop_h = y1 - y0;
    out_img.w = crop_w;
    out_img.h = crop_h;
    out_img.c = in_img.c;
    out_img.w_stride = crop_w * in_img.c;
    out_img.data = (uint8_t*)malloc(crop_h * out_img.w_stride);

    if (!out_img.data) {
        fprintf(stderr, "Memory allocation failed for %d bytes\n", crop_h * out_img.w_stride);
        return;
    }
    
    // 3. 逐行复制像素
    for (int y = 0; y < crop_h; ++y) {
        for (int x = 0; x < crop_w; ++x) {
            int src_pos = (y0 + y) * in_img.w_stride + (x0 + x) * in_img.c;
            int dst_pos = y * out_img.w_stride + x * in_img.c;
            memcpy(out_img.data + dst_pos, in_img.data + src_pos, in_img.c);
        }
    }
}

void draw_boxes(Img &img, const Box_t &box){
    Point pt1 = {
        .x = (int)box.x0,
        .y = (int)box.y0};
    Point pt2 = {
        .x = (int)box.x1,
        .y = (int)box.y1};
    sample_draw_box_for_image(&img, pt1, pt2, color, 2);

}


// 示例用法
int main() {
    FrameInfo input_frame;
    Img input_img;
    std::string image_path = "./imgs/input.png";

    // 读取图片
    read_img(input_frame, input_img, image_path);
    // 检查图片是否成功加载
    if (!input_img.data) {
        fprintf(stderr, "Failed to load input image\n");
        return -1;
    }

    // 定义裁剪区域
    Box_t box = {50, 100, 600, 300};
    // 裁剪图片
    Img crop_img;
    crop_image(input_img, box, crop_img);
    printf("crop_img==> %d %d \n", crop_img.w, crop_img.h);
    // 检查裁剪是否成功，并释放内存
    if (crop_img.data) {
        stbi_write_png("./imgs/cropped.png", crop_img.w, crop_img.h, crop_img.c, 
                      crop_img.data, crop_img.w_stride);
        stbi_image_free(crop_img.data);
    }

    // 缩放图片
    Img resize_img;
    resize_image(crop_img, resize_img, 800, 600);
    // 检查缩放是否成功，并释放内存
    if (resize_img.data) {
        stbi_write_png("./imgs/resized.png", resize_img.w, resize_img.h, 
                      resize_img.c, resize_img.data, resize_img.w_stride);
        stbi_image_free(resize_img.data);  // 必须手动释放
    }

    // yolo缩放填充图片
    Img letterbox_img;
    const uint8_t fill_color[3] = {114, 114, 114}; // YOLO标准填充色
    yolo_resize_img(input_img, letterbox_img, 640, fill_color);
    // 检查缩放是否成功，并释放内存
    if (letterbox_img.data) {
        stbi_write_png("./imgs/letterbox.png", letterbox_img.w, letterbox_img.h, 
                      letterbox_img.c, letterbox_img.data, letterbox_img.w_stride);
        stbi_image_free(letterbox_img.data); // 必须手动释放
    }

    // 绘制检测框
    draw_boxes(input_img, box);
    // 检查绘制是否成功，并释放内存
    if (input_img.data) {
        stbi_write_png("./imgs/drawing.png", input_img.w, input_img.h, input_img.c, 
        input_img.data, input_img.w_stride);
    }
    // 检查绘制是否成功，并释放内存
    if (input_img.data) {
        stbi_image_free(input_img.data);
    }

    return 0;
}